{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 设置DATA_PATH，读取已保存的数据"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ['DATA_PATH'] = 'E:/code/tgtrader/data/akshare_data.db'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "# DuckDBQuery 用于查询duckdb数据库  \n",
    "from tgtrader.utils.duckdb_query import DuckDBQuery\n",
    "# DuckDBQueryDF 用于查询pandas dataframe\n",
    "from tgtrader.utils.duckdb_query_df import DuckDBQueryDF\n",
    "\n",
    "from tgtrader.common import DataSource\n",
    "import time\n",
    "import pandas as pd\n",
    "\n",
    "# ignore warnings\n",
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 设置读取数据库文件的db_query\n",
    "data_source = DataSource.Akshare\n",
    "db_query = DuckDBQuery(data_source)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 数据读取到dataframe"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = db_query.fetch_df(\"select * from t_kdata\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = df.sort_values(by=['code', 'date'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 创建DuckDBQueryDF, 后续可以直接用sql来查询df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "duckdb_df_query = DuckDBQueryDF(df)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 计算SMA"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 用duckdb查询（可直接对df进行查询）"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "duckdb 查询时间: 6.016700299998774 秒\n"
     ]
    }
   ],
   "source": [
    "sql = \"\"\"\n",
    "  SELECT\n",
    "    code,\n",
    "    date,\n",
    "    open,\n",
    "    low,\n",
    "    high,\n",
    "    close,\n",
    "    volume,\n",
    "    -- 计算 20 日 SMA\n",
    "    AVG(close) OVER (PARTITION BY code ORDER BY date ROWS BETWEEN 19 PRECEDING AND CURRENT ROW) AS sma_20\n",
    "  FROM df\n",
    "\"\"\"\n",
    "\n",
    "start_time = time.perf_counter()\n",
    "df_duck = duckdb_df_query.query(sql)\n",
    "end_time = time.perf_counter()\n",
    "print(f\"duckdb 查询时间: {end_time - start_time} 秒\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 用pandas计算"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "pandas 查询时间: 14.132754000002024 秒\n"
     ]
    }
   ],
   "source": [
    "# pandas 查询\n",
    "start_time = time.perf_counter()\n",
    "\n",
    "# 对每个股票计算 sma_20\n",
    "df_pandas = df.groupby('code').apply(lambda x: x.assign(sma_20=x['close'].rolling(window=20, min_periods=0).mean()))\n",
    "\n",
    "end_time = time.perf_counter()\n",
    "print(f\"pandas 查询时间: {end_time - start_time} 秒\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(12108336, 8)"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df_duck.shape\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(12108336, 12)"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df_pandas.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>code</th>\n",
       "      <th>date</th>\n",
       "      <th>open</th>\n",
       "      <th>low</th>\n",
       "      <th>high</th>\n",
       "      <th>close</th>\n",
       "      <th>volume</th>\n",
       "      <th>sma_20</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>3237574</th>\n",
       "      <td>000001</td>\n",
       "      <td>2010-01-04</td>\n",
       "      <td>1402.930054</td>\n",
       "      <td>1353.550049</td>\n",
       "      <td>1406.449951</td>\n",
       "      <td>1355.310059</td>\n",
       "      <td>241923.0</td>\n",
       "      <td>1355.310059</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3237575</th>\n",
       "      <td>000001</td>\n",
       "      <td>2010-01-05</td>\n",
       "      <td>1357.660034</td>\n",
       "      <td>1298.880005</td>\n",
       "      <td>1366.479980</td>\n",
       "      <td>1331.209961</td>\n",
       "      <td>556500.0</td>\n",
       "      <td>1343.260010</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3237576</th>\n",
       "      <td>000001</td>\n",
       "      <td>2010-01-06</td>\n",
       "      <td>1328.270020</td>\n",
       "      <td>1297.119995</td>\n",
       "      <td>1328.270020</td>\n",
       "      <td>1307.699951</td>\n",
       "      <td>412143.0</td>\n",
       "      <td>1331.406657</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3237577</th>\n",
       "      <td>000001</td>\n",
       "      <td>2010-01-07</td>\n",
       "      <td>1307.699951</td>\n",
       "      <td>1278.300049</td>\n",
       "      <td>1316.510010</td>\n",
       "      <td>1293.000000</td>\n",
       "      <td>355337.0</td>\n",
       "      <td>1321.804993</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3237578</th>\n",
       "      <td>000001</td>\n",
       "      <td>2010-01-08</td>\n",
       "      <td>1284.180054</td>\n",
       "      <td>1275.369995</td>\n",
       "      <td>1298.880005</td>\n",
       "      <td>1290.060059</td>\n",
       "      <td>288543.0</td>\n",
       "      <td>1315.456006</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3243190</th>\n",
       "      <td>000001</td>\n",
       "      <td>2024-12-25</td>\n",
       "      <td>2336.639893</td>\n",
       "      <td>2333.389893</td>\n",
       "      <td>2362.639893</td>\n",
       "      <td>2346.389893</td>\n",
       "      <td>1475283.0</td>\n",
       "      <td>2296.251453</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3243191</th>\n",
       "      <td>000001</td>\n",
       "      <td>2024-12-26</td>\n",
       "      <td>2346.389893</td>\n",
       "      <td>2323.639893</td>\n",
       "      <td>2348.020020</td>\n",
       "      <td>2336.639893</td>\n",
       "      <td>1000075.0</td>\n",
       "      <td>2300.476953</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3243192</th>\n",
       "      <td>000001</td>\n",
       "      <td>2024-12-27</td>\n",
       "      <td>2338.270020</td>\n",
       "      <td>2304.129883</td>\n",
       "      <td>2343.139893</td>\n",
       "      <td>2331.760010</td>\n",
       "      <td>1290012.0</td>\n",
       "      <td>2304.133459</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3243193</th>\n",
       "      <td>000001</td>\n",
       "      <td>2024-12-30</td>\n",
       "      <td>2323.639893</td>\n",
       "      <td>2323.639893</td>\n",
       "      <td>2354.520020</td>\n",
       "      <td>2351.270020</td>\n",
       "      <td>1351846.0</td>\n",
       "      <td>2308.684460</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3243194</th>\n",
       "      <td>000001</td>\n",
       "      <td>2024-12-31</td>\n",
       "      <td>2348.020020</td>\n",
       "      <td>2310.639893</td>\n",
       "      <td>2357.770020</td>\n",
       "      <td>2310.639893</td>\n",
       "      <td>1475367.0</td>\n",
       "      <td>2310.391455</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>3573 rows × 8 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "           code        date         open          low         high  \\\n",
       "3237574  000001  2010-01-04  1402.930054  1353.550049  1406.449951   \n",
       "3237575  000001  2010-01-05  1357.660034  1298.880005  1366.479980   \n",
       "3237576  000001  2010-01-06  1328.270020  1297.119995  1328.270020   \n",
       "3237577  000001  2010-01-07  1307.699951  1278.300049  1316.510010   \n",
       "3237578  000001  2010-01-08  1284.180054  1275.369995  1298.880005   \n",
       "...         ...         ...          ...          ...          ...   \n",
       "3243190  000001  2024-12-25  2336.639893  2333.389893  2362.639893   \n",
       "3243191  000001  2024-12-26  2346.389893  2323.639893  2348.020020   \n",
       "3243192  000001  2024-12-27  2338.270020  2304.129883  2343.139893   \n",
       "3243193  000001  2024-12-30  2323.639893  2323.639893  2354.520020   \n",
       "3243194  000001  2024-12-31  2348.020020  2310.639893  2357.770020   \n",
       "\n",
       "               close     volume       sma_20  \n",
       "3237574  1355.310059   241923.0  1355.310059  \n",
       "3237575  1331.209961   556500.0  1343.260010  \n",
       "3237576  1307.699951   412143.0  1331.406657  \n",
       "3237577  1293.000000   355337.0  1321.804993  \n",
       "3237578  1290.060059   288543.0  1315.456006  \n",
       "...              ...        ...          ...  \n",
       "3243190  2346.389893  1475283.0  2296.251453  \n",
       "3243191  2336.639893  1000075.0  2300.476953  \n",
       "3243192  2331.760010  1290012.0  2304.133459  \n",
       "3243193  2351.270020  1351846.0  2308.684460  \n",
       "3243194  2310.639893  1475367.0  2310.391455  \n",
       "\n",
       "[3573 rows x 8 columns]"
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df_duck.query(\"code=='000001'\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>code</th>\n",
       "      <th>date</th>\n",
       "      <th>open</th>\n",
       "      <th>close</th>\n",
       "      <th>high</th>\n",
       "      <th>low</th>\n",
       "      <th>volume</th>\n",
       "      <th>adjust_type</th>\n",
       "      <th>source</th>\n",
       "      <th>create_time</th>\n",
       "      <th>update_time</th>\n",
       "      <th>sma_20</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>10465023</th>\n",
       "      <td>000001</td>\n",
       "      <td>2010-01-04</td>\n",
       "      <td>1402.930054</td>\n",
       "      <td>1355.310059</td>\n",
       "      <td>1406.449951</td>\n",
       "      <td>1353.550049</td>\n",
       "      <td>241923.0</td>\n",
       "      <td>hfq</td>\n",
       "      <td>akshare</td>\n",
       "      <td>1736048554348</td>\n",
       "      <td>1736048554348</td>\n",
       "      <td>1355.310059</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10465024</th>\n",
       "      <td>000001</td>\n",
       "      <td>2010-01-05</td>\n",
       "      <td>1357.660034</td>\n",
       "      <td>1331.209961</td>\n",
       "      <td>1366.479980</td>\n",
       "      <td>1298.880005</td>\n",
       "      <td>556500.0</td>\n",
       "      <td>hfq</td>\n",
       "      <td>akshare</td>\n",
       "      <td>1736048554348</td>\n",
       "      <td>1736048554348</td>\n",
       "      <td>1343.260010</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10465025</th>\n",
       "      <td>000001</td>\n",
       "      <td>2010-01-06</td>\n",
       "      <td>1328.270020</td>\n",
       "      <td>1307.699951</td>\n",
       "      <td>1328.270020</td>\n",
       "      <td>1297.119995</td>\n",
       "      <td>412143.0</td>\n",
       "      <td>hfq</td>\n",
       "      <td>akshare</td>\n",
       "      <td>1736048554348</td>\n",
       "      <td>1736048554348</td>\n",
       "      <td>1331.406657</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10465026</th>\n",
       "      <td>000001</td>\n",
       "      <td>2010-01-07</td>\n",
       "      <td>1307.699951</td>\n",
       "      <td>1293.000000</td>\n",
       "      <td>1316.510010</td>\n",
       "      <td>1278.300049</td>\n",
       "      <td>355337.0</td>\n",
       "      <td>hfq</td>\n",
       "      <td>akshare</td>\n",
       "      <td>1736048554348</td>\n",
       "      <td>1736048554348</td>\n",
       "      <td>1321.804993</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10465027</th>\n",
       "      <td>000001</td>\n",
       "      <td>2010-01-08</td>\n",
       "      <td>1284.180054</td>\n",
       "      <td>1290.060059</td>\n",
       "      <td>1298.880005</td>\n",
       "      <td>1275.369995</td>\n",
       "      <td>288543.0</td>\n",
       "      <td>hfq</td>\n",
       "      <td>akshare</td>\n",
       "      <td>1736048554348</td>\n",
       "      <td>1736048554348</td>\n",
       "      <td>1315.456006</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7769755</th>\n",
       "      <td>000001</td>\n",
       "      <td>2024-12-25</td>\n",
       "      <td>2336.639893</td>\n",
       "      <td>2346.389893</td>\n",
       "      <td>2362.639893</td>\n",
       "      <td>2333.389893</td>\n",
       "      <td>1475283.0</td>\n",
       "      <td>hfq</td>\n",
       "      <td>akshare</td>\n",
       "      <td>1735745107736</td>\n",
       "      <td>1735745107736</td>\n",
       "      <td>2296.251453</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7769756</th>\n",
       "      <td>000001</td>\n",
       "      <td>2024-12-26</td>\n",
       "      <td>2346.389893</td>\n",
       "      <td>2336.639893</td>\n",
       "      <td>2348.020020</td>\n",
       "      <td>2323.639893</td>\n",
       "      <td>1000075.0</td>\n",
       "      <td>hfq</td>\n",
       "      <td>akshare</td>\n",
       "      <td>1735745107736</td>\n",
       "      <td>1735745107736</td>\n",
       "      <td>2300.476953</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7769757</th>\n",
       "      <td>000001</td>\n",
       "      <td>2024-12-27</td>\n",
       "      <td>2338.270020</td>\n",
       "      <td>2331.760010</td>\n",
       "      <td>2343.139893</td>\n",
       "      <td>2304.129883</td>\n",
       "      <td>1290012.0</td>\n",
       "      <td>hfq</td>\n",
       "      <td>akshare</td>\n",
       "      <td>1735745107736</td>\n",
       "      <td>1735745107736</td>\n",
       "      <td>2304.133459</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7769758</th>\n",
       "      <td>000001</td>\n",
       "      <td>2024-12-30</td>\n",
       "      <td>2323.639893</td>\n",
       "      <td>2351.270020</td>\n",
       "      <td>2354.520020</td>\n",
       "      <td>2323.639893</td>\n",
       "      <td>1351846.0</td>\n",
       "      <td>hfq</td>\n",
       "      <td>akshare</td>\n",
       "      <td>1735745107736</td>\n",
       "      <td>1735745107736</td>\n",
       "      <td>2308.684460</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7769759</th>\n",
       "      <td>000001</td>\n",
       "      <td>2024-12-31</td>\n",
       "      <td>2348.020020</td>\n",
       "      <td>2310.639893</td>\n",
       "      <td>2357.770020</td>\n",
       "      <td>2310.639893</td>\n",
       "      <td>1475367.0</td>\n",
       "      <td>hfq</td>\n",
       "      <td>akshare</td>\n",
       "      <td>1735745107736</td>\n",
       "      <td>1735745107736</td>\n",
       "      <td>2310.391455</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>3573 rows × 12 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "            code        date         open        close         high  \\\n",
       "10465023  000001  2010-01-04  1402.930054  1355.310059  1406.449951   \n",
       "10465024  000001  2010-01-05  1357.660034  1331.209961  1366.479980   \n",
       "10465025  000001  2010-01-06  1328.270020  1307.699951  1328.270020   \n",
       "10465026  000001  2010-01-07  1307.699951  1293.000000  1316.510010   \n",
       "10465027  000001  2010-01-08  1284.180054  1290.060059  1298.880005   \n",
       "...          ...         ...          ...          ...          ...   \n",
       "7769755   000001  2024-12-25  2336.639893  2346.389893  2362.639893   \n",
       "7769756   000001  2024-12-26  2346.389893  2336.639893  2348.020020   \n",
       "7769757   000001  2024-12-27  2338.270020  2331.760010  2343.139893   \n",
       "7769758   000001  2024-12-30  2323.639893  2351.270020  2354.520020   \n",
       "7769759   000001  2024-12-31  2348.020020  2310.639893  2357.770020   \n",
       "\n",
       "                  low     volume adjust_type   source    create_time  \\\n",
       "10465023  1353.550049   241923.0         hfq  akshare  1736048554348   \n",
       "10465024  1298.880005   556500.0         hfq  akshare  1736048554348   \n",
       "10465025  1297.119995   412143.0         hfq  akshare  1736048554348   \n",
       "10465026  1278.300049   355337.0         hfq  akshare  1736048554348   \n",
       "10465027  1275.369995   288543.0         hfq  akshare  1736048554348   \n",
       "...               ...        ...         ...      ...            ...   \n",
       "7769755   2333.389893  1475283.0         hfq  akshare  1735745107736   \n",
       "7769756   2323.639893  1000075.0         hfq  akshare  1735745107736   \n",
       "7769757   2304.129883  1290012.0         hfq  akshare  1735745107736   \n",
       "7769758   2323.639893  1351846.0         hfq  akshare  1735745107736   \n",
       "7769759   2310.639893  1475367.0         hfq  akshare  1735745107736   \n",
       "\n",
       "            update_time       sma_20  \n",
       "10465023  1736048554348  1355.310059  \n",
       "10465024  1736048554348  1343.260010  \n",
       "10465025  1736048554348  1331.406657  \n",
       "10465026  1736048554348  1321.804993  \n",
       "10465027  1736048554348  1315.456006  \n",
       "...                 ...          ...  \n",
       "7769755   1735745107736  2296.251453  \n",
       "7769756   1735745107736  2300.476953  \n",
       "7769757   1735745107736  2304.133459  \n",
       "7769758   1735745107736  2308.684460  \n",
       "7769759   1735745107736  2310.391455  \n",
       "\n",
       "[3573 rows x 12 columns]"
      ]
     },
     "execution_count": 39,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df_pandas.query(\"code=='000001'\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
